import traceback
import onnx
from auto_optimizer import OnnxGraph


class Status:
    SUCCESS = 0   # successfully modified the onnx graph
    FAILED = 1   # failed to modify the onnx graph
    NOCHANGE = 2   # onnx graph unchanged after passes


def keep_one_domain(onnx_graph):
    onnx_model = onnx_graph.model()
    is_changed = False
    try:
        for node in onnx_model.graph.node:
            if node.domain != '':
                node.domain = ''
                is_changed = True
        while len(onnx_model.opset_import) > 1:
            onnx_model.opset_import.pop(1)
            is_changed = True
        if is_changed:
            onnx_graph = OnnxGraph.parse(onnx_model)
    except Exception as e:
        print(traceback.print_exc())
        return Status.FAILED, onnx_graph

    status_code = Status.SUCCESS if is_changed else Status.NOCHANGE
    return status_code, onnx_graph


def rename_mod_node(onnx_graph):
    to_rename_nodes = [node for node in onnx_graph.get_nodes('Mod')
                       if '/' in node.name]
    if not to_rename_nodes:
        return Status.NOCHANGE, onnx_graph

    try:
        for node in to_rename_nodes:
            node.name = node.name.replace('/', '_')
    except Exception as e:
        print(traceback.print_exc())
        return Status.FAILED, onnx_graph

    return Status.SUCCESS, onnx_graph
